
#include "pcie.h"
#include "xdma_public.h"
#include <iostream>
#include <string>

using namespace std;

#ifdef _WIN32

#include <Windows.h>
#include <atlbase.h>
#include <atlconv.h>

// ======================= main ===============================================

CPcieComm::CPcieComm() {
    // Initialize critical section mutex.
    InitializeCriticalSection(&cs);

    // PCIE status.
    b_pcie_open = false;
}

CPcieComm::~CPcieComm() {}

/*
 * Function des:
 * In:
 * Out:
 * Return:
 */
BYTE *CPcieComm::allocate_buffer(size_t size, size_t alignment) {
    if (size == 0) {
        size = 4;
    }

    if (alignment == 0) {
        SYSTEM_INFO sys_info;
        GetSystemInfo(&sys_info);
        alignment = sys_info.dwPageSize;
    }
    // verbose_msg("Allocating host-side buffer of size %llu, aligned to %llu
    // bytes\n", size, alignment);
    return (BYTE *)_aligned_malloc(size, alignment);
}

/*
 * Function des:
 * In:
 * Out:
 * Return:
 */
int CPcieComm::get_devices(GUID guid, char *devpath, size_t len_devpath) {
    HDEVINFO device_info = SetupDiGetClassDevs(
        (LPGUID)&guid, NULL, NULL, DIGCF_PRESENT | DIGCF_DEVICEINTERFACE);
    if (device_info == INVALID_HANDLE_VALUE) {
        return -1;
    }

    SP_DEVICE_INTERFACE_DATA device_interface;
    device_interface.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);

    // enumerate through devices
    DWORD index;
    for (index = 0; SetupDiEnumDeviceInterfaces(device_info, NULL, &guid, index,
                                                &device_interface);
         ++index) {

        // get required buffer size
        ULONG detailLength = 0;
        if (!SetupDiGetDeviceInterfaceDetail(device_info, &device_interface, NULL,
                                             0, &detailLength, NULL) &&
            GetLastError() != ERROR_INSUFFICIENT_BUFFER) {
            fprintf(stderr, "SetupDiGetDeviceInterfaceDetail - get length failed\n");
          
            break;
        }

        // allocate space for device interface detail
        PSP_DEVICE_INTERFACE_DETAIL_DATA dev_detail =
            (PSP_DEVICE_INTERFACE_DETAIL_DATA)HeapAlloc(
                GetProcessHeap(), HEAP_ZERO_MEMORY, detailLength);
        if (!dev_detail) {
            fprintf(stderr, "HeapAlloc failed\n");
            break;
        }
        dev_detail->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA);

        // get device interface detail
        if (!SetupDiGetDeviceInterfaceDetail(device_info, &device_interface,
                                             dev_detail, detailLength, NULL,
                                             NULL)) {
            fprintf(stderr, "SetupDiGetDeviceInterfaceDetail - get detail failed\n");
          
            HeapFree(GetProcessHeap(), 0, dev_detail);
            break;
        }
        USES_CONVERSION;
        //StringCchCopyA(devpath, len_devpath, W2A(dev_detail->DevicePath));
        StringCchCopyA(devpath, len_devpath, dev_detail->DevicePath);
        HeapFree(GetProcessHeap(), 0, dev_detail);
    }

    SetupDiDestroyDeviceInfoList(device_info);
    return index;
}

/*
 * Function des:
 * In:
 * Out:
 * Return:
 */
int CPcieComm::cmd_parse(std::vector<std::string> args, BYTE *data, UINT len) {
    Options *option_pt = NULL;

    if (strcmp(args[1].data(), "write") == 0) {
        option_pt = &options_write;
        option_pt->direction = H2C;
    } else if (strcmp(args[1].data(), "read") == 0) {
        option_pt = &options_read;
        option_pt->direction = C2H;
    }

    // device
    option_pt->device = _strdup(args[0].data());
    // base address.
    option_pt->address.QuadPart =
        strtoull(args[2].data(), NULL, 0); // Read address.
    // size
    option_pt->size = len;
    option_pt->alignment = 1;
    // data
    option_pt->data = allocate_buffer(option_pt->size, option_pt->alignment);

    if (option_pt->direction == H2C) // Read operation.
    {
        memcpy(option_pt->data, data, len);
    } else if (option_pt->direction == C2H) // Write operation.
    {
        memset(options_read.data, 0, options_read.size);
    }

    return 1;
}

/*
 * Function des: adjust data sequence.
 * In:
 * Out:
 * Return:
 */
int CPcieComm::data_sequence_adjust(UINT addr, BYTE *data, UINT len,
                                    bool b_write) {
    return 0;
    UINT base = 0;

    if (len != 0) {
        base = addr % len;
    }

    if (b_write) {
        // low bytes to high bytes.
        memcpy(data_buf, data, len);
        memcpy(data, data_buf + base, len - base);
        memcpy(data + (len - base), data_buf, base);
    } else {
        // high bytes to low bytes.
        memcpy(data_buf, data, len);
        memcpy(data + base, data_buf, len - base);
        memcpy(data, data_buf + (len - base), base);
    }

    return 0;
}

//*******************User interface***************************.
/*
 * Function des:
 * In:
 * Out:
 * Return:
 */
int CPcieComm::pcie_open() {
    int status = -1;

    // get device path from GUID
    char device_base_path[MAX_PATH + 1] = "";
    DWORD num_devices = get_devices(GUID_DEVINTERFACE_XDMA, device_base_path,
                                    sizeof(device_base_path));
    printf("Devices found: %d\n", num_devices);
   
    if (num_devices < 1) {
        return status;
    }

    // extend device path to include target device node (xdma_control, xdma_user
    // etc)
    printf("Device base path: %s\n", device_base_path);
    char device_path[MAX_PATH + 1] = "";
    strcpy_s(device_path, sizeof device_path, device_base_path);
    strcat_s(device_path, sizeof device_path, "\\");
    strcat_s(device_path, sizeof device_path, "c2h_0");
    // printf("Device node: %s\n", "c2h_0");

    // open device file
    read_handle = CreateFileA(device_path, GENERIC_READ | GENERIC_WRITE, 0, NULL,
                              OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (read_handle == INVALID_HANDLE_VALUE) {
        fprintf(stderr, "Error opening device, win32 error code: %ld\n",
        GetLastError());
       
        return status;
    }

    // Write channel.
    memset(device_path, 0, sizeof(device_path));
    strcpy_s(device_path, sizeof device_path, device_base_path);
    strcat_s(device_path, sizeof device_path, "\\");
    strcat_s(device_path, sizeof device_path, "h2c_0");
    printf("Device node: %s\n", "h2c_0");

    write_handle = CreateFileA(device_path, GENERIC_READ | GENERIC_WRITE, 0, NULL,
                               OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (write_handle == INVALID_HANDLE_VALUE) {
        fprintf(stderr, "Error opening device, win32 error code: %ld\n",
        GetLastError());
     
        return status;
    }

    // user channel
    memset(device_path, 0, sizeof(device_path));
    strcpy_s(device_path, sizeof device_path, device_base_path);
    strcat_s(device_path, sizeof device_path, "\\");
    //strcat_s(device_path, sizeof device_path, "user");
    strcat_s(device_path, sizeof device_path, "bypass");
    printf("Device node: %s\n", "user");

    user_handle = CreateFileA(device_path, GENERIC_READ | GENERIC_WRITE, 0, NULL,
                               OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (user_handle == INVALID_HANDLE_VALUE) {
        fprintf(stderr, "Error opening device, win32 error code: %ld\n",
        GetLastError());
     
        return status;
    }

    // Update open status.
    b_pcie_open = true;

    return num_devices;
}

/*
 * Function des:
 * In:
 * Out:
 * Return:
 */
int CPcieComm::pcie_write(UINT addr_offset, BYTE *data, UINT len,std::string device) {
    EnterCriticalSection(&cs);

    auto mhandle = device == "user" ? user_handle:write_handle;

    BYTE* data_local = new BYTE[len];

    memcpy(data_local, data, len);

    // Adjust data range.
    //data_sequence_adjust(addr_offset, data_local, len, true);

    char addr[16] = "";

    sprintf(addr, "0x%x", addr_offset);

    // Write data operation.
    std::vector<std::string> argv = {device, "write", addr};

    // parse CMD.
    cmd_parse(argv, data_local, len);

    int retry_count = 0;
    // set file pointer to offset of target address within PCIe BAR
    while (!SetFilePointerEx(mhandle, options_write.address, NULL,
                             FILE_BEGIN) &&
           ++retry_count < 3) {
       
        Sleep(1);
    }

    if (retry_count >= 3) {
        goto Exit_write;
    }

    FlushFileBuffers(mhandle);

    retry_count = 0;
    // Write operation.
    while (!WriteFile(mhandle, options_write.data, options_write.size,
                      &options_write.size, NULL) &&
           ++retry_count < 3) {
       
        Sleep(1);
    }

    if (retry_count >= 3) {
        goto Exit_write;
    }



    FlushFileBuffers(mhandle);

Exit_write:
    if (options_write.data)
    {
        _aligned_free(options_write.data);
        options_write.data = nullptr;
    }

    delete[] data_local;

    LeaveCriticalSection(&cs);

    return retry_count >= 3 ? 0 : options_write.size;
}

/*
 * Function des:
 * In:
 * Out:
 * Return:
 */
int CPcieComm::pcie_read(UINT addr_offset, BYTE *data, UINT &len,std::string device) {
    EnterCriticalSection(&cs);

    auto mhandle = device == "user" ? user_handle:read_handle;

    UINT i = 0;

    char addr[16] = "";
    sprintf(addr, "0x%x", addr_offset);

    // Write data operation.//c2h_0 read 0x0 -l 0x100
    std::vector<std::string> argv = {device, "read", addr};

    memset((void *)&options_read, 0, sizeof(Options));

    // parse CMD.
    cmd_parse(argv, NULL, len);
    // Sleep(10);
    // set file pointer to offset of target address within PCIe BAR
    if (INVALID_SET_FILE_POINTER ==
        SetFilePointerEx(mhandle, options_read.address, NULL, FILE_BEGIN)) {
        fprintf(stderr, "Error setting file pointer, win32 error code: %ld\n",
        GetLastError()); CloseHandle(read_handle);
       
    }

    memset(options_read.data, 0, options_read.size);
    int retry_count = 0;
    while (!ReadFile(mhandle, options_read.data, len, &options_read.size,
                     NULL) &&
           ++retry_count < 3) {
        Sleep(1);
    }

    if (retry_count >= 3) {
    }

    //}
    if ((options_read.data != NULL) && (options_read.size > 0)) {
        // Copy data.
        memcpy(data, options_read.data, options_read.size);
    }

    // Data len
    len = options_read.size;

    if (options_read.data)
        _aligned_free(options_read.data);
    // Adjust data range.
    //data_sequence_adjust(addr_offset, data, len, false);

    LeaveCriticalSection(&cs);

    return options_read.size;
}

/*
 * Function des:
 * In:
 * Out:
 * Return:
 */
bool CPcieComm::pcie_is_open() { return b_pcie_open; }

/*
 * Function des:
 * In:
 * Out:
 * Return:
 */
int CPcieComm::pcie_close() {
    // Check read handle.
    if (read_handle) {
        CloseHandle(read_handle);
    }

    // Check write handle.
    if (write_handle) {
        CloseHandle(write_handle);
    }

    // Update open status.
    b_pcie_open = false;

    return 0;
}
#endif

